#!/usr/bin/env python3
"""
快速诊断过拟合问题
分析训练日志，给出解决建议
"""

import json
import glob
from pathlib import Path
import matplotlib.pyplot as plt

def analyze_training_logs(experiment_name: str):
    """分析训练日志"""
    log_dir = Path(f"codebase/cnclip_finetune/logs/{experiment_name}")
    
    if not log_dir.exists():
        print(f"❌ 日志目录不存在: {log_dir}")
        return
    
    # 查找最新的version
    version_dirs = list(log_dir.glob("version_*"))
    if not version_dirs:
        print(f"❌ 未找到训练日志")
        return
    
    latest_version = max(version_dirs, key=lambda x: x.stat().st_mtime)
    print(f"📂 分析日志: {latest_version}")
    
    # 查找结果文件
    results_dir = Path("research/data") / experiment_name / "results"
    if results_dir.exists():
        result_files = list(results_dir.glob("val_epoch_*.json"))
        if result_files:
            analyze_validation_trend(result_files)
    
    print("\n🎯 建议解决方案:")
    print("1. 立即停止训练")
    print("2. 使用更低学习率 (1e-7 或 5e-8)")
    print("3. 启用早停机制 (patience=3)")
    print("4. 尝试更多冻结层")
    print("5. 减少训练epoch数")

def analyze_validation_trend(result_files):
    """分析验证趋势"""
    epochs = []
    val_r1 = []
    val_loss = []
    
    for file in sorted(result_files):
        try:
            with open(file) as f:
                data = json.load(f)
                epochs.append(data['epoch'])
                val_r1.append(data['retrieval_metrics']['mean_r1'])
                val_loss.append(data['val_loss'])
        except:
            continue
    
    if len(epochs) < 2:
        print("❌ 验证数据不足")
        return
    
    print(f"\n📈 验证趋势分析 ({len(epochs)} 个epoch):")
    print("Epoch | Val R@1  | Val Loss")
    print("-" * 30)
    
    best_r1 = max(val_r1)
    best_epoch = epochs[val_r1.index(best_r1)]
    
    for i, (e, r1, loss) in enumerate(zip(epochs, val_r1, val_loss)):
        marker = "🌟" if r1 == best_r1 else "📉" if i > 0 and r1 < val_r1[i-1] else "📈"
        print(f"{e:5d} | {r1:8.4f} | {loss:8.4f} {marker}")
    
    print(f"\n✨ 最佳性能: Epoch {best_epoch}, R@1={best_r1:.4f}")
    
    # 判断过拟合
    recent_epochs = epochs[-3:] if len(epochs) >= 3 else epochs
    recent_r1 = val_r1[-3:] if len(val_r1) >= 3 else val_r1
    
    if len(recent_r1) >= 2 and all(recent_r1[i] <= recent_r1[i-1] for i in range(1, len(recent_r1))):
        print("🚨 检测到过拟合: 最近几个epoch性能持续下降")
        return True
    
    return False

def suggest_hyperparameters(current_lr: float = 5e-7):
    """建议新的超参数"""
    print(f"\n🔧 超参数调整建议:")
    print(f"当前学习率: {current_lr}")
    print(f"建议学习率: {current_lr/5:.2e} (降低5倍)")
    print(f"备选学习率: {current_lr/10:.2e} (降低10倍)")
    print(f"激进选择: {current_lr/20:.2e} (降低20倍)")
    
    print(f"\n🎛️ 其他建议:")
    print("- 批次大小: 保持64或增加到128")
    print("- 训练轮数: 减少到5-8轮")  
    print("- 早停耐心: 设置为2-3")
    print("- 冻结策略: 尝试 'backbones' 或 'visual+projections'")

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1:
        experiment_name = sys.argv[1]
    else:
        # 尝试找到最新的实验
        log_dir = Path("codebase/cnclip_finetune/logs")
        if log_dir.exists():
            experiments = [d.name for d in log_dir.iterdir() if d.is_dir()]
            if experiments:
                experiment_name = max(experiments, key=lambda x: (log_dir / x).stat().st_mtime)
                print(f"🔍 自动检测最新实验: {experiment_name}")
            else:
                print("❌ 未找到任何实验日志")
                sys.exit(1)
        else:
            print("❌ 日志目录不存在")
            sys.exit(1)
    
    print(f"🚨 过拟合诊断分析")
    print("=" * 50)
    
    analyze_training_logs(experiment_name)
    suggest_hyperparameters()
    
    print("\n" + "=" * 50)
    print("🎯 下一步行动:")
    print("1. 立即停止当前训练 (Ctrl+C)")
    print("2. 保存当前最佳检查点")
    print("3. 使用建议参数重新训练")
    print("4. 监控前几个epoch的验证指标")